import matplotlib.pyplot as plt
import torch
import torch.nn as nn

# 超参数
Epochs = 100
lr = 0.1

# 定义模型
model = nn.Linear(1, 1)

# 损失函数
loss_fn = nn.MSELoss()

# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

''' 学习率调度器 '''
'''
步长衰减调度器StepLR()，在固定epoch更新学习率
优点：实现简单，适用于学习率需要按照固定步长减小的情况；
缺点：不够灵活，步长和衰减率需要手动设置
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer=optimizer,
    step_size=10,
    gamma=0.5       # 衰减率，调整学习率的时候，乘以gamma
)
'''


'''
MultiStepLR()，指定在哪些 epoch 处更新学习率，实际应用中我们时需要观察损失下降曲线的，
然后指定epoch，相比StepLR()更灵活，但是还是需要去指定epoch和衰减率
scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer=optimizer,
    milestones=[20, 40, 60, 80],
    gamma=0.5
)
'''

'''
ExponentialLR() 指数衰减    每个epoch都要衰减一个固定值
优点：实现简单，适用于需要平滑减少学习率的情况；
缺点：衰减率需要手动设置，衰减过程是连续的，不如阶梯式衰减灵活
scheduler = torch.optim.lr_scheduler.ExponentialLR(
    optimizer=optimizer,
    gamma=0.5
)
'''


'''
CosineAnnealingLR() 余弦退火调度器
优点：可以避免局部最小值，周期性的调整有助于探索更多的参数空间；
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer=optimizer,
    T_max=20        # 余弦周期的长度
)
'''

'''
ReduceLROnPlateau() 自动监测训练的性能，若性能改善不佳，则修改学习率
优点：基于模型的实际表现动态调整学习率，适用于不确定何时调整学习率的情况；
缺点：可能需要多个epoch才能调整学习率，反应较慢
'''
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer=optimizer,
    mode="min", # 模式，在"min"模型时，监测的参数不减少的时候，lr要减少
    factor=0.5, # 当检测到性能为改善时，进行学习率调整的系数，类似于gamma
    patience=5,  # 耐心 指的是累计几次未改善，才去动态调整
    threshold=1e-2, # 阈值 前后两次训练损失的差值，小于这个阈值时，则认为训练效果不佳;
    threshold_mode="abs", # abs模式下，min模式下dynamic_threshold = best - threshold
    cooldown=0, # 冷静期
    min_lr=1e-8,    # 能减小到的 最小学习率
    eps=1e-8,       # 防止分母为0的小常数
    verbose=True
)
# 注意注意注意：使用该函数的时候，需要在调用时指定检测的指标(metrics参数) scheduler.step(metrics=loss.item())


# 构造假数据
x_data = torch.tensor([[1], [2], [3], [4]]).float()
y_data = torch.tensor([[2], [4], [6], [8]]).float()

# 循环训练
for epoch in range(Epochs):
    optimizer.zero_grad()
    y_pre = model(x_data)
    loss = loss_fn(y_pre, y_data)
    loss.backward()
    optimizer.step()        # 优化器 更新参数
    scheduler.step(metrics=loss.item())        # 学习率调度器 更新学习率
    # if epoch==0 or (epoch+1)%10==0:
    print(f"[{epoch + 1}/{Epochs}] Loss:{loss.item():.4f} "
              f"LR:{optimizer.param_groups[0]['lr']:.6f}")


# 看一下训练完成之后的，参数
print(f"训练的 w = {model.weight.item():.4f}")
print(f"训练的 b = {model.bias.item():.4f}")















